package lock

import (
	"abdis/lib/utils"
	"sync"
)

// Locks 给 keys 提供读写锁
type Locks struct {
	table []*sync.RWMutex
}

// MakeLocks 创建一个 Locks
func MakeLocks(size int) *Locks {
	table := make([]*sync.RWMutex, size)
	for i := 0; i < size; i++ {
		table[i] = &sync.RWMutex{}
	}
	return &Locks{
		table: table,
	}
}

func (locks *Locks) checkValid() {
	if locks == nil {
		panic("locks is nil")
	}
}

func (locks *Locks) spread(hash uint32) uint32 {
	locks.checkValid()
	size := uint32(len(locks.table))
	return (size - 1) & hash
}

func (locks *Locks) getMutex(key string) *sync.RWMutex {
	index := locks.spread(utils.Fnv32(key))
	return locks.table[index]
}

func (locks *Locks) Lock(key string) {
	locks.getMutex(key).Lock()
}

func (locks *Locks) RLock(key string) {
	locks.getMutex(key).RLock()
}

func (locks *Locks) UnLock(key string) {
	locks.getMutex(key).Unlock()
}

func (locks *Locks) RUnLock(key string) {
	locks.getMutex(key).RUnlock()
}

// Locks 给一组 keys 上锁
func (locks *Locks) Locks(keys ...string) {
	indices := utils.ToLockIndices(keys, false, locks.spread)
	for _, index := range indices {
		m := locks.table[index]
		m.Lock()
	}
}

func (locks *Locks) UnLocks(keys ...string) {
	indices := utils.ToLockIndices(keys, true, locks.spread)
	for _, index := range indices {
		mu := locks.table[index]
		mu.Unlock()
	}
}

func (locks *Locks) RLocks(keys ...string) {
	indices := utils.ToLockIndices(keys, false, locks.spread)
	for _, index := range indices {
		mu := locks.table[index]
		mu.RLock()
	}
}

func (locks *Locks) RUnLocks(keys ...string) {
	indices := utils.ToLockIndices(keys, true, locks.spread)
	for _, index := range indices {
		mu := locks.table[index]
		mu.RUnlock()
	}
}

// RWLocks 同时给读和写的数据上锁
func (locks *Locks) RWLocks(writeKeys []string, readKeys []string) {
	keys := append(writeKeys, readKeys...)
	indices := utils.ToLockIndices(keys, false, locks.spread)
	writeIndexSet := make(map[uint32]struct{})
	for _, wKey := range writeKeys {
		idx := locks.spread(utils.Fnv32(wKey))
		writeIndexSet[idx] = struct{}{}
	}
	for _, index := range indices {
		_, w := writeIndexSet[index]
		m := locks.table[index]
		if w {
			m.Lock()
		} else {
			m.RLock()
		}
	}
}

// RWUnLocks unlocks write keys and read keys together. allow duplicate keys
func (locks *Locks) RWUnLocks(writeKeys []string, readKeys []string) {
	keys := append(writeKeys, readKeys...)
	indices := utils.ToLockIndices(keys, true, locks.spread)
	writeIndexSet := make(map[uint32]struct{})
	for _, wKey := range writeKeys {
		idx := locks.spread(utils.Fnv32(wKey))
		writeIndexSet[idx] = struct{}{}
	}
	for _, index := range indices {
		_, w := writeIndexSet[index]
		mu := locks.table[index]
		if w {
			mu.Unlock()
		} else {
			mu.RUnlock()
		}
	}
}
